package com.sjr.easy.chat.app.ws;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.sjr.easy.chat.app.EasyChatAppConfig;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

@Slf4j
@Component
public class WSConnectionManager implements InitializingBean {

    @Autowired
    private EasyChatAppConfig easyChatAppConfig;

    private final Map<Long, List<WebSocketSession>> wsConnectionMap = Maps.newHashMap();

    private final Lock wsConnectionMapLock = new ReentrantLock();

    public void addSession(Long userId, WebSocketSession session) {
        log.info("add session:{}", userId);
        try {
            wsConnectionMapLock.lock();

            wsConnectionMap.putIfAbsent(userId, Lists.newArrayList());
            List<WebSocketSession> sessionList = wsConnectionMap.get(userId);
            sessionList.add(session);
        } finally {
            wsConnectionMapLock.unlock();
        }
    }

    public void removeSession(Long userId, WebSocketSession session) {
        log.info("remove session:{}", userId);
        try {
            wsConnectionMapLock.lock();

            List<WebSocketSession> sessionList = wsConnectionMap.get(userId);
            if (CollectionUtils.isEmpty(sessionList)) {
                return;
            }
            sessionList.remove(session);
            if (CollectionUtils.isEmpty(sessionList)) {
                wsConnectionMap.remove(userId);
            }
        } finally {
            wsConnectionMapLock.unlock();
        }
    }

    public void sendMessage(Long userId, TextMessage message) {
        try {
            wsConnectionMapLock.lock();
            List<WebSocketSession> sessionList = wsConnectionMap.getOrDefault(userId, Lists.newArrayList());

            for (WebSocketSession webSocketSession : sessionList) {
                try {
                    if (webSocketSession.isOpen()) {
                        webSocketSession.sendMessage(message);
                    }
                } catch (IOException e) {
                    log.error("err:", e);
                }
            }

        } finally {
            wsConnectionMapLock.unlock();
        }
    }

    @Override
    public void afterPropertiesSet() throws Exception {

    }
}
